# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import logging
from typing import List, Union

from tornado import gen

from coin.base.timestamp import get_timestamp
import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex_futures_v3.kr_rest.native_public_client import OkexFuturesNativePublicClient
from coin.exchange.okex_futures_swap.kr_rest.native_public_client import OkexSwapNativePublicClient


class OkexV3RestSubscriberBase(WebsocketJsonSubscriberBase):
  _subscriber_channel = None
  _flow_topic_prefix = None

  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    logger = logger or logging.getLogger(__name__)
    super().__init__(close_callback, logger)
    self._symbol = to_list(symbol)
    if self._subscriber_channel == 'swap':
      # Native symbols are given.
      self._async_feed_client = OkexSwapNativePublicClient(logger=logger)
    elif self._subscriber_channel == 'futures':
      # v1 native symbols are given, cast to v3 native symbols.
      # Before call FromNativeSubscriptionSymbolV1, current_datetime
      # need to be specific to avoid cache mistake.
      # New added USDT futures use v3 native symbols IN UPPER CASE.
      current_datetime = datetime.datetime.now()
      self._symbol = [
          OkexFuturesProduct.FromNativeSubscriptionSymbolV1(
              symbol, current_datetime=current_datetime).native_symbol_v3
          if symbol.islower() else symbol for symbol in self._symbol
      ]
      self._async_feed_client = OkexFuturesNativePublicClient(logger=logger)
    else:
      raise ValueError(self._subscriber_channel)

    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_open_interest = topic_map.search_topic_id(
        '%s_open_interest_%s%s' % (self._flow_topic_prefix, group_id, worker_id))[0]

    self._publisher = {
        'open_interest': queue.get_writer(topic_id=topic_open_interest[1]),
    }

    self._enable_topic_event_publisher(queue, 'okex_ws_topic_event', [topic_open_interest[0]])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_open_interest(self, fut):
    if self.stopped:
      return
    update = fut.result()
    if update is not None:
      message = update.msg
      timestamp = get_timestamp()
      self._publisher['open_interest'].write(timestamp, message)

  @gen.coroutine
  def _run_loop(self):
    try:
      self.on_open()
      while self._state == self.State.RUNNING:
        for symbol in self._symbol:
          fut = self._async_feed_client.async_query_open_interest(symbol)
          if fut is not None:
            self._ioloop.add_future(fut, self.on_open_interest)
        yield gen.sleep(5)
    finally:
      self._notify_close()

  def start(self, ioloop, delay_time=None):
    assert not self.stopped

    self._state = self.State.RUNNING
    self._ioloop = ioloop
    if delay_time is None:
      ioloop.add_callback(self._run_loop)
    else:
      ioloop.call_later(delay_time, self._run_loop)

  def stop(self, trigger_close_callback=True):
    self._async_feed_client.stop()
    super().stop()
